"""
Modbus TCP 工具：将后处理结果打包为浮点寄存器格式，并通过 Modbus TCP 写入寄存器。

本模块实现了以下功能：

1. 将 Python 浮点数转换为大端 16 位寄存器序列。
2. 根据点表或段表生成头部和数据区寄存器列表。
3. 通过 Socket 实现基础的 Modbus TCP 功能码 16（写多个寄存器）写入。

注意：本实现仅用于测试和演示，未完全覆盖所有异常处理。实际接入 PLC 时
建议使用成熟的库（如 `pymodbus`）。
"""

from __future__ import annotations

import logging
import socket
import struct
from typing import Iterable, List, Tuple


def float_to_regs(value: float) -> List[int]:
    """将单个浮点数转换为两个 16 位寄存器（大端）。"""
    # 使用 IEEE754 单精度浮点格式
    packed = struct.pack('>f', value)
    reg1, reg2 = struct.unpack('>HH', packed)
    return [reg1, reg2]


def pack_header(n_records: int, distance: float, status: float = 1.0, version: float = 100.1) -> List[int]:
    """根据协议构建头部寄存器列表。"""
    regs: List[int] = []
    regs += float_to_regs(version)
    regs += float_to_regs(float(n_records))
    regs += float_to_regs(status)
    regs += float_to_regs(distance)
    return regs


def pack_data_points(records: Iterable[List[float]]) -> List[int]:
    """将点表记录转换为寄存器列表。每条记录包含两个浮点数 (flag, z)。"""
    regs: List[int] = []
    for rec in records:
        flag, z = rec
        regs += float_to_regs(flag)
        regs += float_to_regs(z)
    return regs


def pack_data_segments(records: Iterable[List[float]]) -> List[int]:
    """将段表记录转换为寄存器列表。每条记录包含三个浮点数 (flag, z_start, z_end)。"""
    regs: List[int] = []
    for rec in records:
        if len(rec) != 3:
            raise ValueError("段记录必须包含3个浮点数")
        flag, zs, ze = rec
        regs += float_to_regs(flag)
        regs += float_to_regs(zs)
        regs += float_to_regs(ze)
    return regs


def build_payload(records: Iterable[List[float]], distance: float,
                  mode: str = "segments") -> Tuple[List[int], List[int]]:
    """
    根据记录和测距值构建 Modbus 头部和数据区寄存器列表。

    :param records: 点表或段表记录列表。
    :param distance: 相机/刷头到绝缘子的距离，未知可设为 -1。
    :param mode: ``"points"`` 或 ``"segments"``
    :return: (header_regs, data_regs)
    """
    records_list = list(records)
    n_records = len(records_list)
    header_regs = pack_header(n_records, distance, status=1.0)
    if mode.lower() == "points":
        data_regs = pack_data_points(records_list)
    else:
        data_regs = pack_data_segments(records_list)
    return header_regs, data_regs


def _build_modbus_request(transaction_id: int, unit_id: int,
                          function_code: int, start_address: int,
                          regs: List[int]) -> bytes:
    """内部函数：构建 Modbus TCP 写多个寄存器请求帧。"""
    quantity = len(regs)
    # 每个寄存器占两个字节
    byte_count = quantity * 2
    # 构建 PDU
    pdu = struct.pack('>B', unit_id)
    pdu += struct.pack('>B', function_code)
    pdu += struct.pack('>H', start_address)
    pdu += struct.pack('>H', quantity)
    pdu += struct.pack('>B', byte_count)
    # 每个寄存器两个字节，大端
    for reg in regs:
        pdu += struct.pack('>H', reg)
    # MBAP 头：事务标识符、协议标识符、长度
    length = len(pdu)
    header = struct.pack('>HHH', transaction_id, 0x0000, length)
    return header + pdu


def send_modbus_tcp(host: str, port: int, unit_id: int, reg_base: int,
                    header_regs: List[int], data_regs: List[int]) -> None:
    """
    通过 TCP 将寄存器写入 PLC。头部和数据区分别写入。

    写入顺序：先写头部（置 N=0、状态=BUSY），再写数据区，最后写入头部（置实际 N、状态=READY）。
    该函数默认将状态 1.0 视为 BUSY，将状态 2.0 视为 READY。

    :param host: PLC 地址
    :param port: PLC 端口，默认为 502
    :param unit_id: 单元标识符
    :param reg_base: 起始寄存器地址
    :param header_regs: 头部寄存器列表
    :param data_regs: 数据区寄存器列表
    """
    # 分三步写入：BUSY 头、数据、READY 头
    transaction_id = 1
    def _write_registers(regs: List[int], address: int) -> None:
        nonlocal transaction_id
        req = _build_modbus_request(transaction_id, unit_id, 16, address, regs)
        transaction_id = (transaction_id + 1) & 0xFFFF
        with socket.create_connection((host, port), timeout=5.0) as sock:
            sock.sendall(req)
            # 服务器返回响应，读取固定长度即可，不作解析
            _ = sock.recv(1024)
    # 步骤 1：写头部，N=0，状态=BUSY(1)
    busy_header = header_regs.copy()
    # 将 N 和状态改为 BUSY
    # header_regs 格式：[ver_hi, ver_lo, n_hi, n_lo, status_hi, status_lo, dist_hi, dist_lo]
    # 这里修改第 3、4、5、6 个寄存器
    busy_header[2:4] = float_to_regs(0.0)
    busy_header[4:6] = float_to_regs(1.0)
    _write_registers(busy_header, reg_base)
    # 步骤 2：写数据区，从 reg_base + len(header) 开始
    data_address = reg_base + len(header_regs)
    if data_regs:
        _write_registers(data_regs, data_address)
    # 步骤 3：写最终头部，N=实际数量，状态=READY(2)
    ready_header = header_regs.copy()
    # 状态改为 READY
    ready_header[4:6] = float_to_regs(2.0)
    _write_registers(ready_header, reg_base)


__all__ = [
    "float_to_regs", "pack_header", "pack_data_points", "pack_data_segments",
    "build_payload", "send_modbus_tcp"
]